import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import sklearn
import sys
import os
import pandas as pd
import tensorflow as tf
from tensorflow.python import keras
import cv2
import scipy
import PIL


root_directory = "../dataset/明星分类"
height = 128
width = 128
channels = 3
batch_size = 32
num_classes = 2

train_datagen = keras.preprocessing.image.ImageDataGenerator(
    rescale=1./255,
    rotation_range =15,       # 旋转40度
    width_shift_range = 0.2,  # 位移 0.2 个比例 20% 平移
    height_shift_range = 0.2,  #
    shear_range = 0.3,         # 增强强度
    zoom_range = 0.2,
    horizontal_flip = True,
    fill_mode = 'nearest',   # 填充
)
train_generator = train_datagen.flow_from_directory(root_directory,
                                                   target_size = (height,width),
                                                   batch_size = batch_size,
                                                   seed = 7,
                                                   shuffle =True,
                                                   class_mode = "categorical")
validation_datagen = keras.preprocessing.image.ImageDataGenerator(
    rescale=1./255,
)
validation_generator = validation_datagen.flow_from_directory(root_directory,
                                                   target_size = (height,width),
                                                   batch_size = batch_size,
                                                   seed = 7,
                                                   shuffle =False,
                                                   class_mode = "categorical")
train_num = train_generator.samples
valid_num = validation_generator.samples
print(train_num,valid_num)
# 3、开始建立模型
#%%
model = keras.models.Sequential([
    keras.layers.Conv2D(filters=32,kernel_size=3,padding='same',
                        activation='relu',input_shape=[width,height,channels]),
    keras.layers.Conv2D(filters=32,kernel_size=3,padding='same',
                        activation='relu'),
    keras.layers.MaxPool2D(pool_size=2),

    keras.layers.Conv2D(filters=64,kernel_size=3,padding='same',
                        activation='relu'),
    keras.layers.Conv2D(filters=64,kernel_size=3,padding='same',
                        activation='relu'),
    keras.layers.MaxPool2D(pool_size=2),

    keras.layers.Conv2D(filters=64,kernel_size=3,padding='same',
                        activation='relu'),
    keras.layers.Conv2D(filters=64,kernel_size=3,padding='same',
                        activation='relu'),
    keras.layers.MaxPool2D(pool_size=2),
    keras.layers.Flatten(),
    keras.layers.Dense(64, activation='relu'),
    keras.layers.Dense(num_classes, activation='softmax')
])
print(model.summary())
# 加载模型
model.load_weights('_save_weights')
# 开始训练
model.compile(loss='categorical_crossentropy',
              optimizer='adam',
              metrics=['accuracy'])
# epochs = 1
# history = model.fit_generator(train_generator,
#                               steps_per_epoch=train_num // batch_size,
#                               epochs=epochs,
#                               validation_data=validation_generator,
#                               validation_steps=valid_num // batch_size)

# # 打印曲线
# def plot_learing_curves(history,label,epochs,min_value,max_value):
#     data={}
#     data[label] = history.history[label]
#     data['val'+label] = history.history['val_'+label]
#     pd.DataFrame(data).plot(figsize=(8,5))
#     plt.grid(True)
#     plt.axis([0,epochs,min_value,max_value])
#     plt.show()
#
#
# plot_learing_curves(history,'accuracy',epochs,0,1)
# plot_learing_curves(history,'loss',epochs,0,1)

model.save_weights("_save_weights", save_format='tf')  # 保存模型权重
model.save('./model/明星识别.h5') # 保存整个模型
for i in range(1):
    x, y = train_generator.next()
    print(x.shape)
    x1 = x[0]
    x1 = np.reshape(x1,(1,128,128,3))
    y1 = y[0]
    x = x[0]
    x = np.reshape(x,(128,128,3))
    print(x.shape)
    plt.imshow(x)
    plt.show()
    print(model.predict_classes(x1))
    print(y1)

    # score = model.evaluate(x, y, 1)
    # print('Test score:', score[0])
    # print('Test accuracy:', score[1])
    # print('-----------------')
